from visual import *
from random import random
from QMOS import Qmos

class Ballscape(object):
    
    def __init__(self):
        self.color_map = {color.red:"red",
                          color.green:"green",
                          color.blue:"blue",
                          color.yellow:"yellow"}
        self.balls = []
        max_depth = 10
        max_radius = 5
        self.colors = [color.red, color.green, color.blue, color.yellow]
        for i in range(4):
            xpos = max_depth * random()
            ypos = max_depth * random()
            zpos = max_depth * random()
            radius = max_radius * random()
            ball_color = self.colors[i]
            ball = sphere(pos = (xpos, ypos, zpos), radius = radius, color = ball_color)
            self.balls.append(ball)
        
    def oracle(self, index):
        index1 = index & 3
        index2 = index >> 2
        if index1 == index2:
            return False
        ball1 = self.balls[index1]
        ball2 = self.balls[index2]
        pos1 = ball1.pos
        pos2 = ball2.pos
        displacement = pos1 - pos2
        distance = mag(displacement)
        if distance <= ball1.radius + ball2.radius:
            return True
        else:
            return False
            
    def calculate_intersections(self):
        qmos = Qmos(4)
        indices = qmos.quantum_search(self.oracle)
        intersecting_colors = set()
        for index in indices:
            index1 = index & 3
            index2 = index >> 2
            color1 = self.color_map[self.colors[index1]]
            color2 = self.color_map[self.colors[index2]]
            pair = (color1, color2)
            intersecting_colors.add(tuple(sorted(pair)))
        return intersecting_colors

b = Ballscape()
print b.calculate_intersections()